import threading
import time
import cv2
import torch
import robomaster
from robomaster import robot
from robomaster.battery import TelloBattery as Battery
from queue import Queue


"""注意事项：
1. 运行前，请确保公用网络的防火墙已关闭！
"""

# For “NotImplementedError: cannot instantiate 'PosixPath' on your system" bug
import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath
robomaster.config.LOCAL_IP_STR = "192.168.10.2"


class Tello():
    def __init__(self,):
        """Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization."""
        super().__init__()
        # 模型初始化
        self.model = torch.hub.load(r'./', 'custom', path=r"runs/train/exp4/weights/best.pt", source="local")  # load silently)

        # 无人机初始化
        self.drone = robot.Drone()
        self.drone.initialize()
        # 电池初始化
        self.battery = Battery(self.drone)
        # 相机初始化
        self.camera = self.drone.camera
        self.camera.stop_video_stream()
        self.camera.start_video_stream(display=False)
        self.camera.set_fps("high")
        self.camera.set_resolution("high")
        self.camera.set_bitrate(6)

    def infer(self):
        # 使用YOLOv5模型进行推理tello回传的视频流
        img = self.camera.read_cv2_image(strategy="newest")
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 指定视频编码器
        video_writer = cv2.VideoWriter('output.mp4', fourcc, 30.0, (img.shape[1], img.shape[0]))  # 创建视频写入对象
        while True:
            img = self.camera.read_cv2_image(strategy="newest")    #直接读取最新帧，采用="pipeline"会在第二次启动时卡死
            results = self.model(img)     # Inference
            # results.print()
            image = results.show_cv()
            cv2.imshow('results', image)
            # 写入当前帧到视频文件
            video_writer.write(image)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            if not q.empty() and q.get()["close"]:
                break
        video_writer.release()
        cv2.destroyAllWindows()
        print("infer end")

    def get_verson(self):
        # 获取飞机版本信息
        print("Drone sdk version: {0}".format(self.drone.get_sdk_version()))
        # 获取飞机SN信息
        print("drone sn: {0}".format(self.drone.get_sn()))

    def get_battery(self):
        # 获取飞机电量信息
        print("Battery: {}%".format(self.battery.get_battery()))

    def close(self):
        # 结束:释放资源
        self.camera.stop_video_stream()
        self.drone.close()


if __name__ == "__main__":
    q = Queue(100)
    # 初始化
    tello = Tello()
    tello.get_battery()

    # YOLO推理线程
    recv_thread = threading.Thread(target=tello.infer)
    recv_thread.daemon = True
    recv_thread.start()

    # 主线程：用于实现预定动作
    time.sleep(6)
    q.put({"close":True})

    # 结束
    time.sleep(1)      # 延迟1s，确保两个线程都正常结束
    tello.close()
